﻿#ifndef MATRIX_AVX2_H
#define MATRIX_AVX2_H

#include <cmath>
#include <vector>
#include <random>
#include <ostream>
#include <chrono>
#include <cassert>
#include <iostream>
#include <algorithm>
#include <memory>
#include <cstring>
#include <immintrin.h>

#define CUSTOM_ASSERT(condition, message) \
    do { \
        if (!(condition)) { \
            std::cerr << "Assertion failed: " << message << "\n" \
                      << "File: " << __FILE__ << "\n" \
                      << "Line: " << __LINE__ << "\n" \
                      << "Function: " << __func__ << "\n"; \
            std::abort(); \
        } \
    } while (0)


template<typename T>
class BasicMatrix {
public:
    BasicMatrix() : mCols(0), mRows(0) {}

    BasicMatrix(int row, int col)
            : mRows(row), mCols(col), mMat(new T[mCols * mRows], std::default_delete<T[]>()) {
        std::fill_n(mMat.get(), mCols * mRows, T{});
    }

    BasicMatrix(const BasicMatrix<T> &mat)
            : mRows(mat.mRows), mCols(mat.mCols), mMat(new T[mat.mCols * mat.mRows], std::default_delete<T[]>()) {
        std::copy_n(mat.mMat.get(), mCols * mRows, mMat.get());
    }

    BasicMatrix(BasicMatrix<T> &&mat) noexcept
            : mRows(mat.mRows), mCols(mat.mCols), mMat(std::move(mat.mMat)) {
        mat.mRows = 0;
        mat.mCols = 0;
    }

    BasicMatrix &operator=(const BasicMatrix<T> &mat) {
        if (this == &mat) { return *this; }
        mRows = mat.mRows;
        mCols = mat.mCols;
        mMat.reset(new T[mCols * mRows]);
        std::copy_n(mat.mMat.get(), mCols * mRows, mMat.get());
        return *this;
    }

    BasicMatrix &operator=(BasicMatrix<T> &&mat) noexcept {
        if (this == &mat) { return *this; }
        mRows = mat.mRows;
        mCols = mat.mCols;
        mMat = std::move(mat.mMat);

        mat.mRows = 0;
        mat.mCols = 0;
        return *this;
    }


    ~BasicMatrix() = default;

private:
    std::shared_ptr<T[]> mMat;
    size_t mRows, mCols;

public:
    BasicMatrix<T> &operator*=(const BasicMatrix<T> &a) {
        CUSTOM_ASSERT(mCols == a.mRows, "operator*= >>> mCols == a.mRows");
        if (mRows == 0 || mCols == 0 || a.mCols == 0) {
            throw std::invalid_argument("Empty matrix in multiplication");
        }

#ifdef DEBUG
        std::cout << "Multiplying " << mRows << "x" << mCols << " and " << a.mRows << "x" << a.mCols << std::endl;
#endif

        BasicMatrix<T> result(mRows, a.mCols);

        // i-k-j + SIMD
        constexpr size_t SIMD_WIDTH = 4;
#pragma omp parallel for collapse(2)
        for (size_t i = 0; i < mRows; ++i) {
            for (size_t k = 0; k < mCols; ++k) {
                const T r = mMat[i * mCols + k];
                const __m256d r_vec = _mm256_set1_pd(r);

                size_t j = 0;
                for (; j <= a.mCols - SIMD_WIDTH; j += SIMD_WIDTH) {
                    __m256d a_vec = _mm256_load_pd(&a.mMat[k * a.mCols + j]);
                    __m256d res_vec = _mm256_load_pd(&result.mMat[i * a.mCols + j]);
                    res_vec = _mm256_fmadd_pd(r_vec, a_vec, res_vec);
                    _mm256_storeu_pd(&result.mMat[i * a.mCols + j], res_vec);
                }

                for (; j < a.mCols; ++j) {
                    result.mMat[i * a.mCols + j] += r * a.mMat[k * a.mCols + j];
                }
            }
        }

        *this = std::move(result);
        return *this;
    }

    BasicMatrix<T> &operator+=(const BasicMatrix<T> &a) {
        CUSTOM_ASSERT(mRows == a.mRows && mCols == a.mCols, "operator+= >>> mRows == a.mRows && mCols == a.mCols");

        for (size_t i = 0; i < mRows; ++i) {
            for (size_t j = 0; j < mCols; ++j) {
                mMat[i * mCols + j] += a.mMat[i * mCols + j];
            }
        }
        return *this;
    }

    BasicMatrix<T> &operator-=(const BasicMatrix<T> &a) {
        CUSTOM_ASSERT(mRows == a.mRows && mCols == a.mCols, "operator-= >>> mRows == a.mRows && mCols == a.mCols");

        for (size_t i = 0; i < mRows; ++i) {
            for (size_t j = 0; j < mCols; ++j) {
                mMat[i * mCols + j] -= a.mMat[i * mCols + j];
            }
        }
        return *this;
    }


    BasicMatrix<T> &operator*=(double a) {
        static_assert(std::is_same_v<T, double>, "This optimization requires double precision");

        const size_t elements = mRows * mCols;
        constexpr size_t SIMD_WIDTH = 4;

        const __m256d a_vec = _mm256_set1_pd(a);
        double* data = mMat.get();

        size_t i = 0;
        for (; i <= elements - SIMD_WIDTH; i += SIMD_WIDTH) {
            __m256d vec = _mm256_load_pd(&data[i]);
            vec = _mm256_mul_pd(vec, a_vec);
            _mm256_store_pd(&data[i], vec);
        }

        for (; i < elements; ++i) {
            data[i] *= a;
        }
        return *this;
    }


    BasicMatrix<T> &operator/=(double a) {
        CUSTOM_ASSERT(a != 0, "operator/= >>> a != 0");
        for (size_t i = 0; i < mRows; ++i) {
            for (size_t j = 0; j < mCols; ++j) {
                mMat[i * mCols + j] /= a;
            }
        }
        return *this;
    }


    [[nodiscard]] size_t GetRows() const { return mRows; }

    [[nodiscard]] size_t GetCols() const { return mCols; }

    auto &GetData() { return mMat; }

    inline T &At(int row, int col) {
        if (row >= 0 && col >= 0 && row < mRows && col < mCols) {
            return mMat[row * col + col];
        } else {
            throw std::out_of_range("At >>> row >= 0 && col >= 0 && row < mRows && col < mCols");
        }
    }

    void SetSize(int row, int col) {
        if (row == mRows && col == mCols) { return; }
        std::shared_ptr<T[]> newMat(new T[row * col], std::default_delete<T[]>());
        std::fill_n(newMat.get(), row * col, T{});

        int copyRows = row < mRows ? row : mRows;
        int copyCols = col < mCols ? col : mCols;
        for (int i = 0; i < copyRows; ++i) {
            for (int j = 0; j < copyCols; ++j) {
                newMat[i * col + j] = mMat[i * mCols + j];
            }
        }
        mRows = row;
        mCols = col;
        mMat = newMat;
    }

    void Clear() {
        std::memset(mMat.get(), 0, mCols * mRows * sizeof(T));
    }

public:
    template<typename U>
    friend BasicMatrix<U> operator+(const BasicMatrix<U> &a, const BasicMatrix<U> &b);

    template<typename U>
    friend BasicMatrix<U> operator-(const BasicMatrix<U> &a, const BasicMatrix<U> &b);

    template<typename U>
    friend BasicMatrix<U> operator*(const BasicMatrix<U> &a, const BasicMatrix<U> &b);

    template<typename U>
    friend BasicMatrix<U> operator*(const BasicMatrix<U> &a, double b);

    template<typename U>
    friend BasicMatrix<U> operator*(double b, const BasicMatrix<U> &a);

    template<typename U>
    friend BasicMatrix<U> operator/(const BasicMatrix<U> &a, double b);

    template<typename U>
    friend std::ostream &operator<<(std::ostream &out, const BasicMatrix<U> &mat);
};

typedef BasicMatrix<double> Matrix2D;
typedef BasicMatrix<float> Matrix2F;
typedef BasicMatrix<int> Matrix2I;

static void RandomMat(Matrix2D &mat, double a, double b) {
    unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
    std::default_random_engine e(seed);
    std::uniform_real_distribution<double> u(a, b);
    for (int i = 0; i < mat.GetRows(); i++) {
        for (int j = 0; j < mat.GetCols(); j++) {
            mat.At(i, j) = u(e);
        }
    }
}


template<typename U>
BasicMatrix<U> operator+(const BasicMatrix<U> &a, const BasicMatrix<U> &b) {
    CUSTOM_ASSERT(a.mCols == b.mCols && a.mRows == b.mRows, "must a.mCols == b.mCols && a.mRows == b.mRows");

    BasicMatrix<U> out(a.mRows, a.mCols);
    for (int i = 0; i < a.mRows; i++) {
        int tmp_idx = i * a.mCols;
        for (int j = 0; j < a.mCols; j++) {
            out.mMat[tmp_idx + j] = a.mMat[tmp_idx + j] + b.mMat[tmp_idx + j];
        }
    }
    return out;
}

template<typename U>
BasicMatrix<U> operator-(const BasicMatrix<U> &a, const BasicMatrix<U> &b) {
    CUSTOM_ASSERT(a.mCols == b.mCols && a.mRows == b.mRows, "must a.mCols == b.mCols && a.mRows == b.mRows");

    BasicMatrix<U> out(a.mRows, a.mCols);
    for (int i = 0; i < a.mRows; i++) {
        int tmp_idx = i * a.mCols;
        for (int j = 0; j < a.mCols; j++) {
            out.mMat[tmp_idx + j] = a.mMat[tmp_idx + j] - b.mMat[tmp_idx + j];
        }
    }
    return out;
}

template<typename U>
BasicMatrix<U> operator*(const BasicMatrix<U> &a, const BasicMatrix<U> &b) {
    CUSTOM_ASSERT(a.mCols == b.mRows, "must a.mCols == b.mRows");
    BasicMatrix<U> out(a.mRows, b.mCols);

    constexpr size_t SIMD_WIDTH = 4; // 每个AVX2向量处理4个double
    const size_t aRows = a.mRows, aCols = a.mCols, bCols = b.mCols;

#pragma omp parallel for  // OpenMP并行化
    for (size_t i = 0; i < aRows; ++i) {
        for (size_t k = 0; k < aCols; ++k) {
            const U r = a.mMat[i * aCols + k];
            const __m256d r_vec = _mm256_set1_pd(r);

            size_t j = 0;
            for (; j <= bCols - SIMD_WIDTH; j += SIMD_WIDTH) {
                __m256d b_vec = _mm256_load_pd(&b.mMat[k * bCols + j]);
                __m256d out_vec = _mm256_load_pd(&out.mMat[i * bCols + j]);
                out_vec = _mm256_fmadd_pd(r_vec, b_vec, out_vec);
                _mm256_store_pd(&out.mMat[i * bCols + j], out_vec);
            }

            for (; j < bCols; ++j) {
                out.mMat[i * bCols + j] += r * b.mMat[k * bCols + j];
            }
        }
    }

    return out;
}

template<typename U>
BasicMatrix<U> operator*(const BasicMatrix<U> &a, double b) {
    BasicMatrix<U> out(a.mRows, a.mCols);
    for (int i = 0; i < a.mRows; i++) {
        for (int j = 0; j < a.mCols; j++) {
            out.mMat[i * a.mCols + j] = a.mMat[i * a.mCols + j] * b;
        }
    }
    return out;
}

template<typename U>
BasicMatrix<U> operator*(double b, const BasicMatrix<U> &a) {
    return a * b;
}


template<typename U>
BasicMatrix<U> operator/(const BasicMatrix<U> &a, double b) {
    BasicMatrix<U> out(a.mRows, a.mCols);
    for (int i = 0; i < a.mRows; i++) {
        for (int j = 0; j < a.mCols; j++) {
            out.mMat[i * a.mCols + j] = a.mMat[i * a.mCols + j] / b;
        }
    }
    return out;
}


template<typename U>
std::ostream &operator<<(std::ostream &out, const BasicMatrix<U> &mat) {
    for (int i = 0; i < mat.mRows; i++) {
        for (int j = 0; j < mat.mCols; j++) {
            out << mat.mMat[i * mat.mCols + j] << " ";
        }
        out << std::endl;
    }
    return out;
}


#endif //!MATRIX_AVX2_H